import cv2
import os
import time
import numpy as np
import random
import keyboard


KEY_LAYOUT = [
    [keyboard.Q, keyboard.W, keyboard.E, keyboard.R, keyboard.T, keyboard.Y],
    [keyboard.A, keyboard.S, keyboard.D, keyboard.F, keyboard.G, keyboard.H]
]

SKILL_WIDTH = 31
SKILL_HEIGHT = 31
X_GAP = 0
Y_GAP = 0
COLS = 6
ROWS = 2

MATCH_THRESHOLD = 0.1  # 越小越接近，SQDIFF越小越好


def extract_skill_slots(img, output_dir="skills_output", visualize=True):
    # img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError("图像无法加载")

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    os.makedirs(output_dir, exist_ok=True)
    count = 0

    for row in range(ROWS):
        for col in range(COLS):
            key = KEY_LAYOUT[row][col]
            x1 = col * (SKILL_WIDTH + X_GAP)
            y1 = row * (SKILL_HEIGHT + Y_GAP)
            x2 = x1 + SKILL_WIDTH
            y2 = y1 + SKILL_HEIGHT

            crop_color = img[y1:y2, x1:x2]
            crop_gray = gray[y1:y2, x1:x2]

            # 计算平均亮度
            mean_brightness = np.mean(crop_gray)

            # 跳过空图或亮度为 31.13 的技能格
            if crop_color.size == 0 or round(float(mean_brightness), 2) < 35:
                continue

            # 保存图像
            filename = f"{output_dir}/{key}.png"
            cv2.imwrite(filename, crop_color)
            count += 1


def load_templates(template_dir="skills_output"):
    templates = {}
    for row in range(ROWS):
        for col in range(COLS):
            key = KEY_LAYOUT[row][col]
            path = os.path.join(template_dir, f"{key}.png")
            if os.path.exists(path):
                img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
                if img is not None:
                    templates[key] = img
    return templates


def detect_ready_skills(img, templates, visualize=False):
    # img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if img is None:
        raise FileNotFoundError("图像无法加载")

    ready_keys = []

    for row in range(ROWS):
        for col in range(COLS):
            key = KEY_LAYOUT[row][col]
            if key not in templates:
                continue

            x1 = col * (SKILL_WIDTH + X_GAP)
            y1 = row * (SKILL_HEIGHT + Y_GAP)
            x2 = x1 + SKILL_WIDTH
            y2 = y1 + SKILL_HEIGHT
            crop = img[y1:y2, x1:x2]

            if crop.shape != templates[key].shape:
                continue

            res = cv2.matchTemplate(crop, templates[key], cv2.TM_SQDIFF_NORMED)
            score = res[0][0]

            if score < MATCH_THRESHOLD:
                ready_keys.append(key)

    return ready_keys


order = [65, 83, 68, 70, 71, 72,81, 87, 69, 82, 84, 89]


def random_key_selector(image):
    box = (299, 532, 186, 62)
    img = image[box[1]:box[1] + box[3], box[0]:box[0] + box[2]]  # 裁剪图像
    # cv2.imwrite('skill.jpg', img)  # 如需保存裁剪后的图片，取消注释此行
    templates = load_templates("skills_output")

    # 将彩色图像转换为灰度图像
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    keys = detect_ready_skills(img, templates, visualize=False)
    if not keys:
        return keyboard.X
    char_rank = {char: idx for idx, char in enumerate(order)}
    sorted_arr = sorted(keys, key=lambda x: char_rank[x])
    # 获取排序后的第一个元素
    selected_key = sorted_arr[0]
    return selected_key


if __name__ == '__main__':
    img = cv2.imread("skill.jpg")
    random_key_selector(img)